import cv2
import mediapipe as mp
import json
import matplotlib.pyplot as plt
from matplotlib.widgets import Button


def get_pose_data(image_path):
    """
    使用 MediaPipe 提取姿势数据
    """
    mp_pose = mp.solutions.pose
    pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)

    # 读取图片
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Image not found at {image_path}")
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 获取姿势关键点
    results = pose.process(image_rgb)
    if not results.pose_landmarks:
        raise ValueError("No pose landmarks detected!")

    # 提取关键点坐标
    landmarks = [
        {"x": lm.x, "y": lm.y, "z": lm.z, "visibility": lm.visibility}
        for lm in results.pose_landmarks.landmark
    ]
    return landmarks, image.shape


# 保存姿势数据为 JSON
def save_pose_data(landmarks, output_path):
    with open(output_path, "w") as f:
        json.dump(landmarks, f, indent=4)


# 示例调用
image_path = "zishi.png"  # 输入图片路径
landmarks, shape = get_pose_data(image_path)
save_pose_data(landmarks, "./pose_data.json")
print("Pose data saved as JSON.")
